Currently, getting a torch_geometric Dataloader is quite complicated since it requires multiple steps in the process. This PR addresses those concerns by introducing,

1. A method to get torch tensors from getitem call similar to with_format of huggingface datasets Link. This now allows us to get numpy (default), torch or jax arrays from __getitem__ call.

For example,

from openqdc import Dummy

import numpy as np; ds = Dummy(); print (isinstance(ds[0]['positions'], np.ndarray)) # prints True
import torch; ds = Dummy(array_format="torch"); print (isinstance(ds[0]['positions'], torch.Tensor)) # prints True
import jax; ds = Dummy(array_format="jax"); print (isinstance(ds[0]['positions'], jax.numpy.ndarray)) # prints True

2. Added an option to add a transform. In cases, where we want to use the torch_geometric Data object, instead of the sklearn Bunch from getitem, it might be convenient to use a function on the data bunch returned.

from torch_geometric.data import Data
def custom_transform(bunch): 
    return Data(z=bunch.atomic_numbers, pos=bunch.positions, e=bunch.energies, f=bunch.forces)
ds = Dummy(array_format="torch", transform=custom_transform)
print (isinstance(ds[0], Data)) # prints True

3. Last but not the least, in case you want to create a pytorch geometric dataloader, you can do so like this. I still haven't created a method explicitly to do so because the amount of effort now would be reduced quite a bit. If you still think it's valuable for the user to get a dataloader, I can implement it.

Note: I looked through huggingface datasets, they too do not have any such method to get a dataloader from their datasets.

from torch_geometric.data import Data, DataLoader
def custom_transform(bunch): return Data(z=bunch.atomic_numbers, pos=bunch.positions, e=bunch.energies, f=bunch.forces)
ds = Dummy(array_format="torch", transform=custom_transform)
dl = DataLoader(ds, batch_size=4)
batch = next(iter(dl))

TODOs

• Currently I have only added tests for the array_format functionality. Will also add a test for the transform.

@prtos @FNTwin

Checklist:

• Was this PR discussed in a issue? It is recommended to first discuss a new feature into a GitHub issue before opening a PR.
• Add tests to cover the fixed bug(s) or the new introduced feature(s) (if appropriate).
• Update the API documentation is a new function is added or an existing one is deleted.